{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Execute this cell to install dependencies\n",
    "%pip install sf-hamilton[visualization]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Feature Retrieval [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/dagworks-inc/hamilton/blob/main/examples/feast/integration_feature_store/feature_repo/retrieval.ipynb) [![GitHub badge](https://img.shields.io/badge/github-view_source-2b3137?logo=github)](https://github.com/apache/hamilton/blob/main/examples/feast/integration_feature_store/feature_repo/retrieval.ipynb)\n",
    "\n",
    "This notebook shows the basic of feature retrieval using Feast. The `retrieval.ipynb` notebook in `/integration_feature_store/feature_repo` shows more involved usage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The following cell executes the script `run.py` which creates the Feast registry. Ignore the numpy warnings\n",
    "!python run.py\n",
    "\n",
    "from IPython.display import clear_output\n",
    "clear_output()  # clears the jupyter cell output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View the list of feature-views in the Feast registry;\n",
    "# if empty, make sure to execute `run.py` without failure\n",
    "!feast feature-views list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import datetime\n",
    "\n",
    "import feast\n",
    "from hamilton import base, driver\n",
    "\n",
    "import feature_transformations\n",
    "import store_definitions\n",
    "import store_operations  # Feast operations as Hamilton DAG\n",
    "import demo_inputs  # hard-coded inputs for demo\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 1\n",
    "%aimport store_definitions, feature_transformations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visualize the structure of the Feast repository"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dr = driver.Driver(\n",
    "    dict(),\n",
    "    # feature_transformations,\n",
    "    store_definitions,\n",
    "    adapter=base.SimplePythonGraphAdapter(base.DictResult())\n",
    ")\n",
    "\n",
    "dr.display_all_functions(\n",
    "    \"objs\",\n",
    "    render_kwargs={\"format\": \"png\", \"view\": False},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Feast pattern\n",
    "Feast relies on the FeatureStore object to connect to the feature store and interact with it using its methods. The FeatureStore already has a great interface. To use with Hamilton and gain the modularity, reusability, testability, etc. benefits you can usually wrap the FeatureStore in a Hamilton function (more on that later). To learn more about [Feast retrieval](https://docs.feast.dev/getting-started/concepts/feature-retrieval)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# instantiate FeatureStore object\n",
    "feature_store = feast.FeatureStore(repo_path=\".\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Offline / Historical\n",
    "The following cell will retrieve historical data stored in an offline store. Therefore, it will access the Feast DataSource and do the joins at request time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define a retrieval job for entities specified by id and timestamp \n",
    "# `features` specifies columns to retrieve; can be feast.FeatureView, \n",
    "# feast.FeatureService, or column names as string\n",
    "job: feast.infra.offline_stores.offline_store.RetrievalJob = feature_store.get_historical_features(\n",
    "    entity_df=demo_inputs.HISTORICAL_ENTITY_DF,\n",
    "    features=demo_inputs.HISTORICAL_FEATURES,\n",
    ")\n",
    "\n",
    "df = job.to_df()\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Online with materialization\n",
    "The following cells show an online retrieval pattern. First, data from the offline store is materialized for the past `n_days`. This creates the joins and stores the FeatureViews in the online store. Then, features can be easily retrieved."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_days = 90\n",
    "end_date = datetime.datetime.now()\n",
    "start_date = end_date - datetime.timedelta(days=n_days)\n",
    "feature_store.materialize(start_date=start_date, end_date=end_date)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load the FeatureService `driver_activity_v1`\n",
    "service: feast.FeatureService = feature_store.get_feature_service(\"driver_activity_v1\")\n",
    "response: feast.online_response.OnlineResponse = feature_store.get_online_features(\n",
    "    entity_rows=demo_inputs.ONLINE_ENTITY_ROWS,\n",
    "    features=service,\n",
    ")\n",
    "\n",
    "df = response.to_df()\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hamilton pattern\n",
    "With Hamilton, the different operations are defined as function once and stored in the module `store_operations.py`. The variable `final_vars` defines the operations to execute and `inputs` specifies the entities and features to query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dr = driver.Driver(\n",
    "    dict(\n",
    "        feast_repository_path=\".\",\n",
    "        feast_config={},\n",
    "        driver_source_path=\"./data/driver_stats.parquet\",\n",
    "        trips_raw_path=\"./data/trips_raw.parquet\",\n",
    "        trips_stats_3h_path=\"./data/trips_stats_3h.parquet\",\n",
    "    ),\n",
    "    feature_transformations,\n",
    "    store_operations,\n",
    "    store_definitions,\n",
    ")\n",
    "\n",
    "final_vars = [\n",
    "    \"historical_features\",  # specify the function defined in `store_operations`\n",
    "]\n",
    "\n",
    "inputs = dict(\n",
    "    entity_df=demo_inputs.HISTORICAL_ENTITY_DF,\n",
    "    historical_features_=demo_inputs.HISTORICAL_FEATURES,\n",
    "    batch=True,  # triggers the `config.when(batch=True)` for historical retrieval\n",
    ")\n",
    "\n",
    "out = dr.execute(final_vars=final_vars, inputs=inputs)\n",
    "out"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
